#!/usr/bin/env python3
"""修复 ChromaDB 路径配置问题"""

import os
import shutil
from pathlib import Path

def fix_chromadb_path_issue():
    """修复 ChromaDB 路径配置问题"""
    print("=== 修复 ChromaDB 路径配置问题 ===")
    
    # 当前数据所在的正确路径
    correct_path = "./data/chroma_db"
    # ChromaDBClient 默认使用的错误路径
    wrong_path = "./chroma_db"
    
    print(f"正确的数据路径: {correct_path}")
    print(f"代码中使用的路径: {wrong_path}")
    
    # 检查两个路径的情况
    correct_path_exists = os.path.exists(correct_path)
    wrong_path_exists = os.path.exists(wrong_path)
    
    print(f"\n路径检查:")
    print(f"  {correct_path} 存在: {correct_path_exists}")
    print(f"  {wrong_path} 存在: {wrong_path_exists}")
    
    if correct_path_exists and not wrong_path_exists:
        print(f"\n✅ 解决方案: 创建符号链接从 {wrong_path} 指向 {correct_path}")
        try:
            # 在 Windows 上创建目录链接
            if os.name == 'nt':  # Windows
                os.system(f'mklink /D "{wrong_path}" "{correct_path}"')
            else:  # Unix/Linux/Mac
                os.symlink(correct_path, wrong_path)
            print(f"✅ 符号链接创建成功")
        except Exception as e:
            print(f"❌ 创建符号链接失败: {e}")
            print(f"💡 备选方案: 复制数据到 {wrong_path}")
            try:
                shutil.copytree(correct_path, wrong_path)
                print(f"✅ 数据复制成功")
            except Exception as e2:
                print(f"❌ 数据复制也失败: {e2}")
                return False
    
    elif wrong_path_exists and not correct_path_exists:
        print(f"\n✅ 解决方案: 移动数据从 {wrong_path} 到 {correct_path}")
        try:
            os.makedirs(os.path.dirname(correct_path), exist_ok=True)
            shutil.move(wrong_path, correct_path)
            print(f"✅ 数据移动成功")
        except Exception as e:
            print(f"❌ 数据移动失败: {e}")
            return False
    
    elif correct_path_exists and wrong_path_exists:
        print(f"\n⚠️  两个路径都存在，需要手动处理")
        print(f"建议: 检查两个路径的数据，决定保留哪个")
        return False
    
    else:
        print(f"\n❌ 两个路径都不存在，没有数据可以修复")
        return False
    
    return True

def verify_fix():
    """验证修复结果"""
    print(f"\n=== 验证修复结果 ===")
    
    import chromadb
    
    # 测试使用错误路径（代码中的默认路径）
    wrong_path = "./chroma_db"
    try:
        client = chromadb.PersistentClient(path=wrong_path)
        collections = client.list_collections()
        print(f"✅ 使用路径 {wrong_path} 成功连接")
        print(f"   找到 {len(collections)} 个集合")
        
        # 检查 conversations 集合
        for collection in collections:
            if collection.name == "conversations":
                count = collection.count()
                print(f"   conversations 集合包含 {count} 个文档")
                
                if count > 0:
                    # 测试搜索
                    results = collection.query(
                        query_texts=["王者荣耀"],
                        n_results=3
                    )
                    print(f"   搜索 '王者荣耀' 找到 {len(results['ids'][0])} 个结果")
                    
                    # 测试带用户过滤的搜索
                    filtered_results = collection.query(
                        query_texts=["王者荣耀"],
                        n_results=3,
                        where={"user_id": "default_user"}
                    )
                    print(f"   搜索 '王者荣耀' (用户: default_user) 找到 {len(filtered_results['ids'][0])} 个结果")
                break
        
        return True
        
    except Exception as e:
        print(f"❌ 验证失败: {e}")
        return False

if __name__ == "__main__":
    success = fix_chromadb_path_issue()
    if success:
        verify_fix()
    else:
        print("\n❌ 修复失败，请手动检查路径配置")